package calcite;

import org.apache.calcite.adapter.jdbc.JdbcSchema;
import org.apache.calcite.interpreter.Bindables;
import org.apache.calcite.jdbc.CalciteConnection;
import org.apache.calcite.plan.RelOptTable;
import org.apache.calcite.plan.hep.HepPlanner;
import org.apache.calcite.plan.hep.HepProgram;
import org.apache.calcite.rel.RelHomogeneousShuttle;
import org.apache.calcite.rel.RelNode;
import org.apache.calcite.rel.RelShuttle;
import org.apache.calcite.rel.core.TableScan;
import org.apache.calcite.rel.logical.LogicalTableScan;
import org.apache.calcite.schema.SchemaPlus;
import org.apache.calcite.tools.FrameworkConfig;
import org.apache.calcite.tools.Frameworks;
import org.apache.calcite.tools.RelBuilder;
import org.apache.calcite.tools.RelRunner;

import javax.sql.DataSource;
import java.sql.*;
import java.util.concurrent.atomic.AtomicInteger;

public class CalciteJdbcExample {

    private static final String POSTGRESQL_SCHEMA = "PUBLIC";
    private static final String MYSQL_SCHEMA = "mysql";

    public static void main(String[] args) throws Exception {

        // Build our connection
        Connection connection = DriverManager.getConnection("jdbc:calcite:");

        // Unwrap our connection using the CalciteConnection
        CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);

        // Get a pointer to our root schema for our Calcite Connection
        SchemaPlus rootSchema = calciteConnection.getRootSchema();

        // Instantiate a data source, this can be autowired in using Spring as well
        DataSource postgresDataSource = JdbcSchema.dataSource(
                "jdbc:postgresql://localhost/db",
                "org.postgresql.Driver", // Change this if you want to use something like MySQL, Oracle, etc.
                "postgres", // username
                "example"   // password
        );

        // Instantiate a data source, this can be autowired in using Spring as well
        DataSource mysqlDataSource = JdbcSchema.dataSource(
                "jdbc:mysql://localhost/db",
                "com.mysql.jdbc.Driver", // Change this if you want to use something like MySQL, Oracle, etc.
                "Username", // username
                "Password"   // password
        );

        // Attach our Postgres Jdbc Datasource to our Root Schema
        rootSchema.add(POSTGRESQL_SCHEMA, JdbcSchema.create(rootSchema, POSTGRESQL_SCHEMA, postgresDataSource, null, null));

        // Attach our MySQL Jdbc Datasource to our Root Schema
        rootSchema.add(MYSQL_SCHEMA, JdbcSchema.create(rootSchema, MYSQL_SCHEMA, mysqlDataSource, null, null));


        // Build a framework config to attach to our Calcite Planners and  Optimizers
        FrameworkConfig config = Frameworks.newConfigBuilder()
                .defaultSchema(rootSchema)
                .build();

        RelBuilder rb = RelBuilder.create(config);

        RelNode node = rb
                // First parameter is the Schema, the second is the table name
                .scan("PUBLIC", "TABLE_NAME_IN_POSTGRES")
                .scan("mysql", "TABLE_NAME_IN_MYSQL")
                // If you want to select from more than one table, you can do so by adding a second scan parameter
                .filter(
                        rb.equals(rb.field("fieldname"), rb.literal("literal"))
                )
                // These are the fields you want to return from your query
                .project(
                        rb.field("id"),
                        rb.field("col1"),
                        rb.field("colb")
                )
                .build();


        HepProgram program = HepProgram.builder().build();
        HepPlanner planner = new HepPlanner(program);

        planner.setRoot(node);

        RelNode optimizedNode = planner.findBestExp();

        final RelShuttle shuttle = new RelHomogeneousShuttle() {
            @Override public RelNode visit(TableScan scan) {
                final RelOptTable table = scan.getTable();
                if (scan instanceof LogicalTableScan && Bindables.BindableTableScan.canHandle(table)) {
                    return Bindables.BindableTableScan.create(scan.getCluster(), table);
                }
                return super.visit(scan);
            }
        };

        optimizedNode = optimizedNode.accept(shuttle);

        final RelRunner runner = connection.unwrap(RelRunner.class);
        PreparedStatement ps = runner.prepareStatement(optimizedNode);

        ps.execute();

        ResultSet resultSet = ps.getResultSet();
       // DBTablePrinter.printResultSet(resultSet);
    }

    private void init(){
        String jdbcUrl = "" ;
        DataSource mockDataSource = JdbcSchema.dataSource(jdbcUrl, "org.hsqldb.jdbcDriver", "", "");
    }


    public void calciteUnitTestExample() throws SQLException, SQLException {

        // Build our connection
        Connection connection = DriverManager.getConnection("jdbc:calcite:");

        // Unwrap our connection using the CalciteConnection
        CalciteConnection calciteConnection = connection.unwrap(CalciteConnection.class);

        // Get a pointer to our root schema for our Calcite Connection
        SchemaPlus rootSchema = calciteConnection.getRootSchema();

        // Attach our Postgres Jdbc Datasource to our Root Schema
       // rootSchema.add("exampleSchema", JdbcSchema.create(rootSchema, "exampleSchema", mockDataSource, null, null));

        FrameworkConfig config = Frameworks.newConfigBuilder()
                .defaultSchema(rootSchema)
                .build();

        RelBuilder r = RelBuilder.create(config);

        RelNode node = r
                // First parameter is the Schema, the second is the table name
                .scan("exampleSchema", "MYTEMPTABLE")
                // If you want to select from more than one table, you can do so by adding a second scan parameter
                .filter(
                        r.equals(r.field("ID"), r.literal(1))
                )
                // These are the fields you want to return from your query
                .project(
                        r.field("ID"),
                        r.field("EXAMPLEFOO"),
                        r.field("EXAMPLEBAR")
                )
                .build();

        RelRunner runner = connection.unwrap(RelRunner.class);
        PreparedStatement ps = runner.prepareStatement(node);

        ps.execute();

        ResultSet resultSet = ps.getResultSet();
        //DBTablePrinter.printResultSet(resultSet);
    }

    static class MockDb {
        MockDb() {}
        static final MockDb INSTANCE = new MockDb();
        private final AtomicInteger id = new AtomicInteger(1);

        public String getUrl() {
            return "jdbc:hsqldb:mem:db" + id.getAndIncrement();
        }
    }
}